package elementbyid

import (
	"golang.org/x/net/html"
)

func forEachNode(n *html.Node, pre, post func(n *html.Node) bool) {
	if pre != nil && !pre(n) {
		return
	}
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		forEachNode(c, pre, post)
	}

	if post != nil && !post(n) {
		return
	}
}

func ElementByID(doc *html.Node, id string) *html.Node {
	var foundNode *html.Node
	forEachNode(doc, func(n *html.Node) bool {
		if foundNode != nil {
			return false
		}
		if n.Type == html.ElementNode && hasID(n, id) {
			foundNode = n
		}
		return true
	},
		func(n *html.Node) bool {
			b := foundNode == nil
			return b
		})
	return foundNode
}

func hasID(n *html.Node, id string) bool {
	for _, attr := range n.Attr {
		if attr.Key == "id" && attr.Val == id {
			return true
		}
	}
	return false
}

func Extend(s string, f func(string) string) string {
	return f(s)
}
